import sympy as sym
import numpy as np

def any(A, dim=1):
    """
    Tests whether at least one element of A returns logical 1 (true). 
    If A is a matrix, any tests elements of each column. 
    If A is a multidimensional array, any tests elements along one dimension.

    Args:
        A : symbolic array
        dim : Dimension(1: test each column, 2: test each row).

    Returns:
        the result of a and b(a & b).
    """
    if isinstance(A, (list, np.ndarray)):
        narr = np.array(A)
        result = None
        # Symbolic Vection
        if len(narr.shape) == 1:
            result = False
            for expr in narr:
                result |= expr
            return result
        # Along the column, result is a ndarray with shape(1,col)
        if dim == 1:
            # print(narr.shape)
            result = np.full(np.shape(np.sum(np.zeros(np.shape(narr)), axis=0)), False)
            for row in range(narr.shape[0]):
                for col in range(narr.shape[1]):
                    result[col] |= narr[row][col]
            return result
        # Along the row, result is a ndarray with shape(row,1)
        elif dim == 2:
            # print(narr.shape)
            result = np.full(np.shape(np.sum(np.zeros(np.shape(narr)), axis=1)), False)
            # print(result)
            for row in range(narr.shape[0]):
                for col in range(narr.shape[1]):
                    result[row] |= narr[row][col]
            # Reshape it to a column vector
            return result.reshape((narr.shape[0],-1))
        else:
            raise ValueError("Illegal Dim argument.")
    else:
       raise TypeError("Input type not supported.") 


# Test case 1: Input is symbolic expressions
x = sym.symbols('x')
arr = [sym.Ne(x, x+1), sym.Ge(abs(x) ,0) , sym.Eq(x , x)]
assert any(arr) == True

# Test case 2: Input is 2 dim symbolic expressions with dim 1 (default)
x = sym.symbols('x')
arr = [
    [sym.Ne(x ,x+1), sym.Ge(abs(x) ,0) , sym.Eq(x ,x+1)],
    [sym.Eq(x ,x+1), sym.Lt(abs(x) ,0) , sym.Ne(x ,x)]
    ]
result = np.array([True,True,False])
assert np.array_equal(any(arr,1), result)

# Test case 3: Input is 2 dim symbolic expressions with dim 2
x = sym.symbols('x')
arr = [
    [sym.Ne(x ,x+1), sym.Ge(abs(x) ,0) , sym.Eq(x ,x)],
    [sym.Eq(x ,x+1), sym.Lt(abs(x) ,0) , sym.Ne(x ,x)]
    ]
result = np.array([[True],[False]])
assert np.array_equal(any(arr,2), result)